#include "parser.hpp"
#include "ast.hpp"
#include "visitor.hpp"
#include "symtable.hpp"
#include "typechecker.hpp"

#include <stdlib.h>
#include <iostream>
using namespace std;
#include <cassert>

Parser::Parser(Lexer* lex, SymbolTable* symtable) {
    _lex = lex;
    _cur_token = NULL;
    _sym_table = symtable;
}

Parser::~Parser() {
    if (_lex) {
        delete _lex;
        _lex = NULL;
    }

    if (_cur_token) {
        delete _cur_token;
        _cur_token = NULL;
    }
}

Token* Parser::get_token() {
    if (_cur_token == NULL) {
        _cur_token = _lex->next();
        if (_cur_token)
            _cur_token->print();
    }
    
    return _cur_token;
}

void Parser::consume() {
    _cur_token = NULL;
}

void Parser::match(TokenType tt) {
    get_token();
    if (_cur_token->_tt != tt) {
        printf("Parser Error: expected %d, but got %d\n", tt, _cur_token->_tt);
        assert(_cur_token->_tt == tt); // to crash.
    }
    else {
        consume();
    }
}

Node* Parser::stmts() {
    ListNode* l = new ListNode();
    Token* t = get_token();
    while (t != NULL && t->_tt != T_RIGHT_BRACE) {
        l->add(stmt());
        t = get_token();
    }

    return l;
}

Node* Parser::stmt() {
    Token* t = get_token();
    if (t->_tt == T_KW_PRINT) {
        Node* n = print_stmt();
        t = get_token();
        match(T_ENDL);
        return n;
    }
    else if (t->_tt == T_KW_RETURN) {
        Node* n = return_stmt();
        t = get_token();
        match(T_ENDL);
        return n;
    }
    else if (t->_tt == T_KW_IF) {
        return if_stmt();
    }
    else if (t->_tt == T_KW_WHILE) {
        return while_stmt();
    }
    else if (t->_tt == T_KW_VAR) {
        Node* n = var_def_stmt();
        t = get_token();
        match(T_ENDL);
        return n;
    }
    else if (t->_tt == T_KW_VAL) {
        Node* n = val_def_stmt();
        t = get_token();
        match(T_ENDL);
        return n;
    }
    else if (t->_tt == T_NAME) {
        Node* n = expression();
        t = get_token();
        if (t->_tt == T_ASN) {
            n = assign_stmt(n);
        }
        t = get_token();
        match(T_ENDL);
        return n;
    }
    else if (t->_tt == T_KW_DEF) {
        return func_def();
    }
    else if (t->_tt == T_KW_TYPE) {
        Node* n = typedef_stmt();
        t = get_token();
        match(T_ENDL);
        return n;
    }

    return NULL;
}

Node* Parser::if_stmt() {
    match(T_KW_IF);
    match(T_LEFT_PAR);
    Node* cond = or_test();
    match(T_RIGHT_PAR);
    Node* then_part = code_block();

    Token* t = get_token();
    Node* else_part = NULL;
    if (t != NULL && t->_tt == T_KW_ELSE) {
        match(T_KW_ELSE);
        else_part = code_block();
    }

    return new IfNode(cond, then_part, else_part);
}

Node* Parser::while_stmt() {
    match(T_KW_WHILE);
    match(T_LEFT_PAR);
    Node* cond = or_test();
    match(T_RIGHT_PAR);
    Node* loop_body = code_block();

    return new WhileNode(cond, loop_body);
}

Node* Parser::type_atom() {
    Token* t = get_token();

    if (t->_tt == T_NAME) {
        Node* n = new TypeVarNode(string(t->_value, t->_length));
        consume();
        return n;
    }
    else if (t->_tt == T_INT) {
        Node* n = new TypeConstInt(stoi(t));
        consume();
        return n;
    }
    else if (t->_tt == T_LEFT_PAR) {
        consume();
        Node* n = type_stmt();
        match(T_RIGHT_PAR);
        return n;
    }

    return NULL;
}

/*
 * type_stmt : type_atom "->" type_stmt
 *              | type_atom type_trails
 */
Node* Parser::type_stmt() {
    Node* a = type_atom();

    Token* t = get_token();
    if (t->_tt == T_ARROW) {
        consume();
        a = new ArrowTypeNode(a, type_stmt());
    }
    else if (t->_tt == T_LEFT_BRACKET) {
        consume();
        a = new TypeCallNode(a, type_arguments());
        match(T_RIGHT_BRACKET);
    }

    return a;
}

Node* Parser::val_def_stmt() {
    Node* val_type = NULL;
    Node* init_value = NULL;
    
    match(T_KW_VAL);
    Token* t = get_token();
    if (t->_tt != T_NAME) {
        printf("Parse Error: expect name but got %d\n", t->_tt);
        exit(-1);
    }

    string name(t->_value, t->_length);
    consume();

    t = get_token();
    if (t->_tt == T_COLON) {
        match(T_COLON);
        val_type = type_stmt();
    }

    t = get_token();
    match(T_ASN);
    init_value = expression();

    VarDefNode* n = new VarDefNode(name, val_type, init_value);
    TypeChecker(_sym_table, false).visit(n);

    return n;
}

Node* Parser::var_def_stmt() {
    Node* var_type = NULL;
    Node* init_value = NULL;
    
    match(T_KW_VAR);
    Token* t = get_token();
    if (t->_tt != T_NAME) {
        printf("Parse Error: expect name but got %d\n", t->_tt);
        exit(-1);
    }

    char* name = new char[t->_length + 1];
    memcpy(name, t->_value, t->_length);
    name[t->_length] = 0;

    consume();
    t = get_token();

    if (t->_tt == T_COLON) {
        match(T_COLON);
        var_type = type_stmt();
    }

    t = get_token();
    if (t->_tt == T_ASN) {
        match(T_ASN);
        init_value = expression();
    }

    if (var_type == NULL && init_value == NULL) {
        printf("Parse Error: variable must be declared with type or init value.\n");
        exit(-1);
    }

    VarDefNode* n = new VarDefNode(name, var_type, init_value);
    TypeChecker(_sym_table, false).visit(n);

    return n;
}

Node* Parser::assign_stmt(Node* left) {
    match(T_ASN);
    return new AssignNode(left, expression());
}

Node* Parser::new_clause() {
    match(T_KW_NEW);
    Node* n = type_stmt();
    Token* t = get_token();
    assert(t->_tt == T_LEFT_PAR);
    return trails(n);
}

Node* Parser::param_def() {
    Node* var_type = NULL;
    Node* init_value = NULL;;
    
    match(T_KW_VAR);
    Token* t = get_token();
    if (t->_tt != T_NAME) {
        printf("Parse Error: expect name but got %d\n", t->_tt);
        exit(-1);
    }

    string name(t->_value, t->_length);

    consume();
    t = get_token();
    if (t->_tt == T_COLON) {
        match(T_COLON);
        var_type = type_stmt();
    }

    t = get_token();
    if (t->_tt == T_ASN) {
        match(T_ASN);
        init_value = expression();
    }

    if (var_type == NULL && init_value == NULL) {
        printf("Parse Error: variable must be declared with type or init value.\n");
        exit(-1);
    }

    return new ParamDefNode(name, var_type, init_value);
}

ListNode* Parser::parameters() {
    Token* t = get_token();
    ListNode* n = new ListNode();

    while (t->_tt != T_RIGHT_PAR) {
        n->add(param_def());
        t = get_token();
        if (t->_tt == T_COMMA) {
            match(T_COMMA);
        }
        t = get_token();
    }

    return n;
}

Node* Parser::func_def() {
    match(T_KW_DEF);
    Token* t = get_token();
    FuncDefNode* n = new FuncDefNode(t);
    match(T_NAME);

    // generic check
    t = get_token();
    if (t->_tt == T_LEFT_BRACKET) {
        consume();
        n->set_generic_args(type_parameters());
        match(T_RIGHT_BRACKET);
    }

    match(T_LEFT_PAR);
    n->set_args(parameters());
    match(T_RIGHT_PAR);

    t = get_token();
    if (t->_tt == T_COLON) {
        match(T_COLON);
        n->set_ret_type(type_stmt());
    }

    n->set_body(code_block());
    TypeChecker(_sym_table, false).fill_sym_table(n);

    return n;
}

Node* Parser::code_block() {
    match(T_LEFT_BRACE);
    Node* n = stmts();
    match(T_RIGHT_BRACE);
    return n;
}

CmpNode* Parser::cmp_op() {
    Token* t = get_token();
    CmpNode* n = NULL;

    if (t->_tt == T_LT) {
        match(T_LT);
        n = new CmpNode(CmpNode::CMP_LT);   
    }
    else if (t->_tt == T_EQU) {
        match(T_EQU);
        n = new CmpNode(CmpNode::CMP_EQU);
    }
    else if (t->_tt == T_LE) {
        match(T_LE);
        n = new CmpNode(CmpNode::CMP_LE);
    }

    return n;
}

Node* Parser::print_stmt() {
    match(T_KW_PRINT);
    match(T_LEFT_PAR);
    Node* n = new PrintNode(expression());
    match(T_RIGHT_PAR);

    return n;
}

Node* Parser::return_stmt() {
    match(T_KW_RETURN);
    return new ReturnNode(expression());
}

/*
 * lambda_def = "$" NAME ":" type "=>" expression
 *            | "$" NAME ":" type "=>" suite
 */
Node* Parser::lambda_def() {
    match(TokenType::T_LAMBDA);
    Node* n = param_def();
    match(TokenType::T_LAM_DEF);
    ListNode* codes = new ListNode();
    codes->add(new ReturnNode(expression()));

    return new LambdaDef(n, codes);
}

/*
 * expression -> lambda_def
 *              | or_test
 *              | or_test IF or_test ELSE or_test
 */
Node* Parser::expression() {
    Token* left_most = get_token();
    if (left_most->_tt == TokenType::T_LAMBDA) {
        return lambda_def();
    }

    Node* a = or_test();

    return a;
}

/*
 * or_test -> and_test
 *          | or_test OR and_test
 */
Node* Parser::or_test() {
    Node* a = and_test();
    Token* op = get_token();
    while (op != NULL && op->_tt == TokenType::LOGIC_OR) {
        match(LOGIC_OR);
        Node* b = and_test();
        a = new BinaryOp(AST_OP_LOG_OR, a, b);
        op = get_token();
    }

    return a;
}

/*
 * and_test -> not_test
 *          | not_test "&&" and_test
 */
Node* Parser::and_test() {
    Node* a = not_test();
    Token* op = get_token();
    while (op != NULL && op->_tt == TokenType::LOGIC_AND) {
        consume();
        Node* b = and_test();
        a = new BinaryOp(AST_OP_LOG_AND, a, b);
        op = get_token();
    }

    return a;
}

/*
 * not_test -> "!" not_test
 *           | comparison
 */
Node* Parser::not_test() {
    Token* t = get_token();

    if (t->_tt == TokenType::T_LOGIC_NOT) {
        consume();
        return new LogicNotNode(not_test());
    }

    return comparison();
}

Node* Parser::comparison() {
    Node* a = or_expr();

    CmpNode* cmp = cmp_op();
    while (cmp != NULL) {
        Node* b = or_expr();
        cmp->set_left(a);
        cmp->set_right(b);
        a = cmp;
        cmp = cmp_op();
    }

    return a;
}

Node* Parser::or_expr() {
    Node* a = xor_expr();
    Token* op = get_token();
    while (op != NULL && op->_tt == TokenType::BIT_OR) {
        consume();
        Node* b = xor_expr();
        a = new BinaryOp(AST_OP_BIT_OR, a, b);
        op = get_token();
    }

    return a;
}

Node* Parser::xor_expr() {
    Node* a = and_expr();
    Token* op = get_token();
    while (op != NULL && op->_tt == TokenType::BIT_XOR) {
        consume();
        Node* b = and_expr();
        a = new BinaryOp(AST_OP_BIT_XOR, a, b);
        op = get_token();
    }

    return a;
}

Node* Parser::and_expr() {
    Node* a = shift_expr();
    Token* op = get_token();
    while (op != NULL && op->_tt == TokenType::BIT_AND) {
        consume();
        Node* b = shift_expr();
        a = new BinaryOp(AST_OP_BIT_AND, a, b);
        op = get_token();
    }

    return a;
}

Node* Parser::shift_expr() {
    Node* a = arith_expr();
    Token* op = get_token();
    while (op != NULL &&
            (op->_tt == TokenType::LEFT_SHIFT ||
             op->_tt == TokenType::RIGHT_SHIFT)) {
        TokenType op_tt = op->_tt;
        consume();
        Node* b = arith_expr();
        if (op_tt == TokenType::LEFT_SHIFT) {
            a = new BinaryOp(AST_OP_LEFT_SHIFT, a, b);
        }
        else {
            a = new BinaryOp(AST_OP_RIGHT_SHIFT, a, b);
        }
        op = get_token();
    }

    return a;
}

Node* Parser::arith_expr() {
    Node* a = term();
    Token* op = get_token();
    while (op != NULL &&
        (op->_tt == TokenType::T_PLUS || op->_tt == TokenType::T_MINUS)) {
        TokenType op_tt = op->_tt;
        consume();
        Node* b = term();
        if (op_tt == TokenType::T_PLUS) {
            a = new BinaryOp(AST_OP_ADD, a, b);
        } else {
            a = new BinaryOp(AST_OP_SUB, a, b);
        }

        op = get_token();
    }

    return a;
}

Node* Parser::term() {
    Node* a = factor();
    Token* op = get_token();
    while (op != NULL &&
        (op->_tt == T_MULT || op->_tt == T_DIV)) {
        consume();
        Node* b = factor();
        if (op->_tt == T_MULT) {
            a = new BinaryOp(AST_OP_MUL, a, b);
        } else {
            a = new BinaryOp(AST_OP_DIV, a, b);
        }
        
        op = get_token();
    }
    
    return a;
}

bool Parser::is_trailer_char(TokenType tt) {
    return (tt == TokenType::T_LEFT_PAR) ||
        (tt == TokenType::T_LEFT_BRACKET);
}

Node* Parser::trails(Node* left) {
    Token* t = get_token();
    Node* n = nullptr;

    // call a function
    if (t->_tt == TokenType::T_LEFT_PAR) {
        consume();
        ListNode* args = expr_list();
        match(TokenType::T_RIGHT_PAR);
        n = new CallNode(left, args);
    }
    // subscript
    else if (t->_tt == TokenType::T_LEFT_BRACKET) {
        consume();
        TypeChecker tc(_sym_table, false);
        tc.visit(left);
        Node* a = nullptr;
        if (tc.expr_type()->is_generic_func()) {
            a = type_arguments();
        }
        else {
            a = expression();
        }
        match(TokenType::T_RIGHT_BRACKET);
        n = new SubscrNode(left, a);
    }

    t = get_token();
    while (is_trailer_char(t->_tt)) {
        n = trails(n);
        t = get_token();
    }

    return n;
}

Node* Parser::atom() {
    Token* data = get_token();
    if (data->_tt == T_INT) {
        consume();
        return new ConstInt(stoi(data));
    }
    else if (data->_tt == T_STRING) {
        consume();
        return new ConstString(data->_length, data->_value);
    }
    else if (data->_tt == T_CONST_UNIT) {
        consume();
        return BuiltinConst::UnitNode;
    }
    else if (data->_tt == T_CONST_TRUE) {
        consume();
        return new VarNode("True");
    }
    else if (data->_tt == T_CONST_FALSE) {
        consume();
        return new VarNode("False");
    }
    else if (data->_tt == T_NAME) {
        VarNode* n = new VarNode(data);
        consume();
        return n;
    }
    else if (data->_tt == T_LEFT_PAR) {
        match(T_LEFT_PAR);
        Node* a = expression();
        match(T_RIGHT_PAR);
        
        return a;
    }
    else if (data->_tt == T_LEFT_BRACKET) {
        match(T_LEFT_BRACKET);
        Node* a = new ConstArray(expr_list());
        match(T_RIGHT_BRACKET);
        return a;
    }
    else if (data->_tt == T_KW_NEW) {
        return new_clause();
    }
    return NULL;
}

ListNode* Parser::expr_list() {
    ListNode* n = new ListNode();
    Token* t = get_token();

    while ((t->_tt != T_RIGHT_BRACKET)
        && (t->_tt != T_RIGHT_PAR)) {
        n->add(expression());
        t = get_token();
        if (t->_tt == T_COMMA) {
            match(T_COMMA);
            t = get_token();
            continue;
        }
        else if (t->_tt == T_RIGHT_BRACKET) {
            break;
        }
        else if (t->_tt == T_RIGHT_PAR) {
            break;
        }
        else {
            match(T_COMMA); // error
        }
    }

    return n;
}

/*
 * factor -> "-" factor
 *          | atom
 *          | atom trails
 */
Node* Parser::factor() {
    Token* data = get_token();
    if (data->_tt == T_MINUS) {
        consume();
        return new BinaryOp(AST_OP_SUB, new ConstInt(0), factor());
    }

    Node* n = atom();
    Token* t = get_token();

    if (is_trailer_char(t->_tt)) {
        return trails(n);
    }
    else {
        return n;
    }
}

Node* Parser::type_param() {
    Token* t = get_token();
    if (t->_tt == T_NAME) {
        Node* n = new TypeVarNode(string(t->_value, t->_length));
        consume();
        return n;
    }
    else {
        match(T_NAME);
        return nullptr;
    }
}

ListNode* Parser::type_arguments() {
    Token* t = get_token();
    ListNode* n = new ListNode();

    if (t->_tt == T_RIGHT_BRACKET) {
        return n;
    }

    n->add(type_stmt());
    t = get_token();

    while (t->_tt == T_COMMA) {
        consume();
        n->add(type_stmt());
        t = get_token();
    }

    return n;
}

/*
 * type_parameters : type_param [, type_param]*
 */
ListNode* Parser::type_parameters() {
    Token* t = get_token();
    ListNode* n = new ListNode();

    n->add(type_param());
    t = get_token();

    while (t->_tt == T_COMMA) {
        consume();
        n->add(type_param());
        t = get_token();
    }

    return n;
}

/**
 * typedef_stmt -> "type" NAME "=" type_stmt
 *              | "type" NAME "[" poly_var_list "]" "=" type_stmt
 */
Node* Parser::typedef_stmt() {
    match(TokenType::T_KW_TYPE);
    Token* t = get_token();
    assert(t->_tt == TokenType::T_NAME);
    string name(t->_value, t->_length);
    match(TokenType::T_NAME);

    t = get_token();
    if (t->_tt == T_LEFT_BRACKET) {
        consume();
        TypeFuncDefNode* n = new TypeFuncDefNode(name);
        n->set_args(type_parameters());
        match(TokenType::T_RIGHT_BRACKET);
        match(TokenType::T_ASN);
        n->set_body(type_stmt());
        return n;
    }
    else {
        match(TokenType::T_ASN);
        TypeVarDefNode* n = new TypeVarDefNode(name, type_stmt());
        return n;
    }

    return nullptr;
}

int Parser::stoi(Token* data) {
    int value = 0;
    for (int i = 0; i < data->_length; i++) {
        value = value * 10 + data->_value[i] - '0';
    }
    
    return value;
}

Node* Parser::parse() {
    return stmts();
}

